import cv2
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from util import try_gpu

class MyDataset(Dataset):
    def __init__(self, path, is_train=True):
        self.path = path
        if is_train:
            txt_name = "/train.txt"
        else:
            txt_name = "/test.txt"
        file = open(self.path + txt_name, 'r')
        base_data = []
        for line in file:
            line = line.rstrip('\n')
            line = line.split()
            img_o_path = line[0]
            img_m_path = line[1]
            label = line[2]
            base_data.append((img_o_path, img_m_path, label))
        self.data = base_data
        
    def __getitem__(self, index):
        img_o_path, img_m_path, label = self.data[index]
        img_o = cv2.imread(img_o_path)
        img_m = cv2.imread(img_m_path)
        img = cv2.absdiff(img_o, img_m)
        img = torch.tensor(img, dtype=float).to(torch.float32)
        label = torch.tensor(int(label))
        img = img.permute(2, 0, 1)
        return img, label

    def __len__(self):
        return len(self.data)

def get_dataload(path, my_batch_size, is_train=True):
    dataset = MyDataset(path, is_train)
    dataloader = DataLoader(dataset, batch_size=my_batch_size, num_workers=24, shuffle=True)
    return dataloader
